/**
 * SHACrypt command line
 *
 * SHACrypt - Encrypt files using cryptographic hash functions
 * Copyright (C) 2010 LlamaSlayers
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * @since 1.1.2
 * @version 1.2.3
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <time.h>
#include "shacrypt.h"
#include "../config.h"

#ifdef WINDOWS
#include <windows.h>

/**
 * Adapted from http://www.dreamincode.net/code/snippet2415.htm
 *
 * This function will simulate the getch() function from conio.h.
 * Note that this does not work for special keys like F1-F12, arrow keys,
 * even the ESC key does not seem to work with this function.
 *
 * @since 1.2.3
 * @version 1.2.3
 */
unsigned char getch() {
	DWORD mode;
	TCHAR theChar = '\0';
	DWORD count;
	//get a handle to stdin
	HANDLE hStdIn = GetStdHandle(STD_INPUT_HANDLE);
	//save the current input mode...
	GetConsoleMode(hStdIn,&mode);
	// Clear the mode, turn off ENABLE_ECHO_INPUT and ENABLE_LINE_INPUT
	// so that the output will not be echoed and will not pause until the end of
	// a line for input.
	SetConsoleMode(hStdIn,mode & ~ENABLE_ECHO_INPUT & ~ENABLE_LINE_INPUT);
    WaitForSingleObject( hStdIn, INFINITE );
	// Read in 1 char from the input buffer.
	ReadConsoleA(hStdIn,&theChar,sizeof(TCHAR),&count, NULL);
	//restore the current input mode.
	SetConsoleMode(hStdIn, mode);
	return (unsigned char)theChar;
}
#else
#include <termios.h>
#include <unistd.h>

/**
 * From http://cboard.cprogramming.com/faq-board/27714-faq-there-getch-conio-equivalent-linux-unix.html
 *
 * @since 1.2.3
 * @version 1.2.3
 */
unsigned char getch() {
    struct termios oldt,
                   newt;
    unsigned char ch;
    tcgetattr( STDIN_FILENO, &oldt );
    newt = oldt;
    newt.c_lflag &= ~( ICANON | ECHO );
    tcsetattr( STDIN_FILENO, TCSANOW, &newt );
    ch = getchar();
    tcsetattr( STDIN_FILENO, TCSANOW, &oldt );
    return ch;
}
#endif

static const char identifier[4] = "\xA7\x09\xC3";

/**
 * Display the version number and copyright/licensing
 * in a GNU-standard way.
 *
 * @since 1.2.2
 * @version 1.2.2
 */
void showversion() {
    printf( "%s (%s) %s\n"
            "\n"
            "Copyright (C) 2010 LlamaSlayers\n"
            "License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>\n"
            "This is free software: you are free to change and redistribute it.\n"
            "There is NO WARRANTY, to the extent permitted by law.\n",
            PACKAGE_TARNAME, PACKAGE_NAME, PACKAGE_VERSION );
}

/**
 * Handle all help-related commands
 *
 * @since 1.1.2
 * @version 1.2.3
 * @todo Make this use GNU style.
 */
void showhelp( char *appname, char *mode ) {
    if ( strcmp( mode, "ha" ) == 0 || strcmp( mode, "hasymmetric" ) == 0 ) {
        printf( "SHACrypt %s: Asymmetric Encryption\n", PACKAGE_VERSION );
        printf( "Please note: All steps are formed in the same way as normal encryption, with the exception of the first argument (the e, a, or d). See '%s h' for more details. In these examples, Alice is the sender and Bob is the recipient. Between each step, the file would be transmitted to the other person.\n", appname );
        printf( " Step 1, Alice: Encryption - '%s e'\n", appname );
        printf( " Step 2, Bob: Asymmetric encryption - '%s a'\n", appname );
        printf( " Step 3, Alice: Asymmetric decryption - '%s a'\n", appname );
        printf( " Step 4, Bob: Decryption - '%s d'\n", appname );
    } else if ( strcmp( mode, "he" ) == 0 || strcmp( mode, "hencryption" ) == 0 ) {
        printf( "SHACrypt %s: Advanced encryption\n", PACKAGE_VERSION );
        printf( "Appending letters to the \"%s e\" will change the hashing algorithms used.\n", appname );
        printf( " a - SHA1\n" );
        printf( " b - SHA224\n" );
        printf( " c - SHA256\n" );
        printf( " d - SHA384\n" );
        printf( " e - SHA512\n" );
        printf( " f - MD5\n" );
        printf( " g - CubeHash\n" );
        printf( " h - WHIRLPOOL\n" );
        printf( "Multiple letters can be combined. If none are used, SHACrypt will default to SHA1.\n" );
    } else if ( strcmp( mode, "h" ) == 0 || strcmp( mode, "help" ) == 0 || strcmp( mode, "hhelp" ) == 0 ) {
        printf( "Usage: %s MODE[FLAGS] [PASSWORD] [INFILE OUTFILE]\n"
                "Securely encrypt files using cryptographic hash functions.\n"
                "\n"
                " hTOPIC, --help[=TOPIC]:      display help\n"
                " Help topics:\n"
                "     asymmetric (a),\n"
                "     encryption (e)\n"
                "\n"
                " v, --version:                display version imformation\n"
                "\n"
                " eTYPES, --encrypt[=TYPES]:   encrypt a file using TYPES hash algorithms.\n"
                "                              See %s --help=encryption for more details\n"
                " d, --decrypt:                decrypt a file\n"
                " c, --compat:                 encrypt or decrypt a SHACrypt 1.0 file\n"
                " a, --asymmetric:             switch the password on an encrypted file\n"
                "                              See %s --help=asymmetric for more details\n"
                "\n"
                "Report bugs to <%s> or to the project site at <%s>.\n",
                PACKAGE, PACKAGE, PACKAGE, PACKAGE_BUGREPORT, PACKAGE_URL );
    } else {
        fprintf( stderr, "%s: unknown option: %s\n"
                         "See %s --help for help.\n",
                         PACKAGE, mode, appname );
    }
}

/**
 * Encrypt the file named in argv[3] to the file named in argv[4]
 * using the password in argv[2]. If there are less than 3 argv,
 * use stdin and stdout.
 *
 * @since 1.0.1
 * @version 1.2.3
 */
void encryptmode( int argc, char **argv ) {
    int n;
    FILE *f,
         *o;
    unsigned char fbuf[256] = {0},
                  ibuf[256] = {0},
                  obuf[256] = {0},
                  randkey[256] = {0},
                  scramblekey[256] = {0},
                  password[2048] = {0},
                  algos = '\0';
    SHACrypt_Context s;

    if ( strlen( argv[1] ) > 1 ) {
        for ( n = 1; n < strlen( argv[1] ); n++ ) {
            switch ( argv[1][n] ) {
                case 'a': // SHA1
                    algos |= '\x01';
                    break;
                case 'b': // SHA224
                    algos |= '\x02';
                    break;
                case 'c': // SHA256
                    algos |= '\x04';
                    break;
                case 'd': // SHA384
                    algos |= '\x08';
                    break;
                case 'e': // SHA512
                    algos |= '\x10';
                    break;
                case 'f': // MD5
                    algos |= '\x20';
                    break;
                case 'g': // CubeHash
                    algos |= '\x40';
                    break;
                case 'h': // WHIRLPOOL
                    algos |= '\x80';
                    break;
                default:
                    fprintf( stderr, "%s: unknown switch: %c\n"
                                     "See %s --help=encryption for help.\n",
                                     PACKAGE, argv[1][n], argv[0] );
                    exit( EXIT_FAILURE );
            }
        }
    }

    if ( algos == '\0' )
        algos = '\x01';

    if ( argc == 4 ) {
        f = fopen( argv[2], "rb");
        if ( !f ) {
            perror( argv[2] );
            exit( EXIT_FAILURE );
        }
        o = fopen( argv[3], "wb" );
        if ( !o ) {
            perror( argv[3] );
            exit( EXIT_FAILURE );
        }
    } else if ( argc > 4 ) {
        f = fopen( argv[3], "rb");
        if ( !f ) {
            perror( argv[3] );
            exit( EXIT_FAILURE );
        }
        o = fopen( argv[4], "wb" );
        if ( !o ) {
            perror( argv[4] );
            exit( EXIT_FAILURE );
        }
    } else {
        f = stdin;
        o = stdout;
    }

    n = SHACrypt_GetRand( algos, randkey );

    if ( argc == 2 ) {
        printf( "You must specify a password or two filenames or both.\n" );
        exit( EXIT_FAILURE );
    } else if ( argc == 4 ) {
        printf( "Enter Password: " );
        unsigned char ch;
        int i = 0;
        do {
            ch = getch();
            if ( ch == '\b' ) {
                if ( i > 0 ) {
                    printf( "\b \b" );
                    i--;
                }
            } else {
                password[i] = ch;
                printf( "*" );
                i++;
            }
        } while ( ch != '\n' && ch != '\r' && i < 2048 );
        printf( "\b \b\n" );
        password[i - 1] = '\0';
    } else {
        strcpy( password, argv[2] );
    }

    SHACrypt_Init( &s, password, strlen( password ), algos, randkey, n );

    fwrite( identifier, 1, 3, o );        // Identifier
    fwrite( "\3\0", 1, 2, o );            // Format version
    fputc( (unsigned char)( n - 1 ), o ); // Random seed length
    fwrite( randkey, 1, n, o );           // Random seed
    fputc( algos, o );                    // Hash alogithms used

    n = SHACrypt_GetRand( algos, scramblekey );

    /* We need to also encrypt the scramble key */
    fbuf[0] = (unsigned char)( n - 1 );
    SHACrypt_Process( &s, fbuf, obuf, 1 );
    fputc( obuf[0], o );                 // Scramble key length
    SHACrypt_Process( &s, scramblekey, obuf, n );
    fwrite( obuf, 1, n, o );             // Scramble key

    SHACrypt_SetScramble( &s, scramblekey, n );

    /* Do the actual encryption */
    while ( ( n = fread( fbuf, 1, 256, f ) ) > 0 ) {
        SHACrypt_Scramble( &s, fbuf, ibuf, n );
        SHACrypt_Process( &s, ibuf, obuf, n );

        fwrite( obuf, 1, n, o );
    }
}

void decryptmode( int argc, char **argv ) {
    int n, version;
    FILE *f,
         *o;
    unsigned char fbuf[256] = {0},
                  ibuf[256] = {0},
                  obuf[256] = {0},
                  randkey[256] = {0},
                  scramblekey[256] = {0},
                  algos,
                  randlen,
                  scramblelen;
    SHACrypt_Context s;

    if ( argc > 3 ) {
        f = fopen( argv[3], "rb");
        if ( !f ) {
            perror( argv[3] );
            exit( EXIT_FAILURE );
        }
        o = fopen( argv[4], "wb" );
        if ( !o ) {
            perror( argv[4] );
            exit( EXIT_FAILURE );
        }
    } else {
        f = stdin;
        o = stdout;
    }

    fread( fbuf, 1, 5, f );

    if ( strncmp( fbuf, identifier, 3 ) != 0 ) {
        fprintf( stderr, "SHACrypt file identifier not found. This is not a SHACrypt file." );
        exit( EXIT_FAILURE );
    }
    if ( fbuf[3] == '\1' && fbuf[4] == '\0' ) {
        randlen = '\x13';
        fread( randkey, 1, 20, f );
        algos = '\x01';
        version = 2;
    } else if ( fbuf[3] == '\2' && fbuf[4] == '\0' ) {
        randlen = fgetc( f );
        fread( randkey, 1, (int)randlen + 1, f );
        algos = fgetc( f );
        version = 2;
    } else if ( fbuf[3] == '\3' && fbuf[4] == '\0' ) {
        randlen = fgetc( f );
        fread( randkey, 1, (int)randlen + 1, f );
        algos = fgetc( f );
        version = 3;
    } else {
        fprintf( stderr, "Input file was not in a format that could be read by SHACrypt. (Maybe it's a newer version?)" );
        exit( EXIT_FAILURE );
    }

    SHACrypt_Init( &s, argv[2], strlen( argv[2] ), algos, randkey, (int)randlen + 1 );

    /* Get the scramble key */
    if ( version >= 3 ) {
        fbuf[0] = fgetc( f );
        SHACrypt_Process( &s, fbuf, obuf, 1 );
        scramblelen = obuf[0];

        n = fread( fbuf, 1, (int)scramblelen + 1, f );
        SHACrypt_Process( &s, fbuf, scramblekey, n );

        SHACrypt_SetScramble( &s, scramblekey, n );
    }

    /* Do the actual decryption */
    while ( ( n = fread( fbuf, 1, 256, f ) ) > 0 ) {
        if ( version >= 3 ) {
            SHACrypt_Process( &s, fbuf, ibuf, n );
            SHACrypt_Unscramble( &s, ibuf, obuf, n );
        } else {
            SHACrypt_Process( &s, fbuf, obuf, n );
        }

        fwrite( obuf, 1, n, o );
    }
}

void compatmode( int argc, char **argv ) {
    int n;
    FILE* f;
    FILE* o;
    static char fbuf[256] = {0},
                obuf[256] = {0};
    SHACrypt_Context s;

    if ( argc > 3 ) {
        f = fopen( argv[3], "rb");
        if ( !f ) {
            perror( argv[3] );
            exit( EXIT_FAILURE );
        }
        o = fopen( argv[4], "wb" );
        if ( !o ) {
            perror( argv[4] );
            exit( EXIT_FAILURE );
        }
    } else {
        f = stdin;
        o = stdout;
    }

    SHACrypt_Init( &s, argv[2], strlen( argv[2] ), '\x01', "", 0 );

    /* Do the actual en/decryption */
    while ( ( n = fread( fbuf, 1, 256, f ) ) > 0 ) {
        SHACrypt_Process( &s, fbuf, obuf, n );

        fwrite( obuf, 1, n, o );
    }
}

void asymmetricmode( int argc, char **argv ) {
    int n;
    FILE *f,
         *o;
    SHACrypt_Context s;
    unsigned char algos,
                  randlen,
                  *version,
                  fbuf[256] = {0},
                  obuf[256] = {0};

    if ( argc > 3 ) {
        f = fopen( argv[3], "rb");
        if ( !f ) {
            perror( argv[3] );
            exit( EXIT_FAILURE );
        }
        o = fopen( argv[4], "wb" );
        if ( !o ) {
            perror( argv[4] );
            exit( EXIT_FAILURE );
        }
    } else {
        f = stdin;
        o = stdout;
    }

    fread( fbuf, 1, 5, f );


    if ( strncmp( fbuf, identifier, 3 ) != 0 ) {
        fprintf( stderr, "SHACrypt file identifier not found. This is not a SHACrypt file." );
        exit( EXIT_FAILURE );
    }
    if ( fbuf[3] == '\1' && fbuf[4] == '\0' ) {
        randlen = '\x13';
        fread( fbuf, 1, 20, f );
        algos = '\x01';
        version = "\2\0";
    } else if ( fbuf[3] == '\2' && fbuf[4] == '\0' ) {
        randlen = fgetc( f );
        fread( fbuf, 1, (int)randlen + 1, f );
        algos = fgetc( f );
        version = "\2\0";
    } else if ( fbuf[3] == '\3' && fbuf[4] == '\0' ) {
        randlen = fgetc( f );
        fread( fbuf, 1, (int)randlen + 1, f );
        algos = fgetc( f );
        version = "\3\0";
    } else {
        fprintf( stderr, "Input file was not in a format that could be read by SHACrypt. (Maybe it's a newer version?)" );
        exit( EXIT_FAILURE );
    }

    fwrite( identifier, 1, 3, o );           // Identifier
    fwrite( version, 1, 2, o );              // Format version
    fputc( randlen, o );                     // Random seed length
    fwrite( fbuf, 1, (int)randlen + 1, o ); // Random seed
    fputc( algos, o );                       // Hash algorithms used

    SHACrypt_Init( &s, argv[2], strlen( argv[2] ), algos, fbuf, (int)randlen + 1 );

    /* Do the actual encryption */
    while ( ( n = fread( fbuf, 1, 256, f ) ) > 0 ) {
        SHACrypt_Process( &s, fbuf, obuf, n );

        fwrite( obuf, 1, n, o );
    }
}

/**
 * 
 *
 * @since 1.0
 * @version 1.2.2
 */
int main( int argc, char **argv ) {
    // Remove leading - from the command unless it has two
    if ( argc > 1 && argv[1][0] == '-' && argv[1][1] != '-' )
        argv[1]++;

    /**
     * Help
     */
    if ( argc == 1 || strcmp( argv[1], "--help" ) == 0 ) {
        showhelp( argv[0], "h" );
    } else if ( argv[1][0] == 'h' ) {
        showhelp( argv[0], argv[1] );
    } else if ( strncmp( argv[1], "--help=", 7 ) == 0 ) {
        argv[1] += 6;
        argv[1][0] = 'h';
        showhelp( argv[0], argv[1] );
    } else
    /**
     * Version
     */
    if ( argv[1][0] == 'v' || strcmp( argv[1], "--version" ) == 0 ) {
        showversion();
    } else
    /**
     * Compat mode
     */
    if ( argv[1][0] == 'c' || strcmp( argv[1], "--compat" ) == 0 ) {
        compatmode( argc, argv ); // compatmode doesn't use any flags, so we can just pass it in.
    } else
    /**
     * Encrypt mode
     */
    if ( strcmp( argv[1], "--encrypt" ) == 0 ) {
        argv[1] = "e";
        encryptmode( argc, argv );
    } else if ( argv[1][0] == 'e' ) {
        encryptmode( argc, argv );
    } else if ( strncmp( argv[1], "--encrypt=", 10 ) == 0 ) {
        argv[1] += 9;
        argv[1][0] = 'e';
        encryptmode( argc, argv );
    } else
    /**
     * Decrypt mode
     */
    if ( argv[1][0] == 'd' || strcmp( argv[1], "--decrypt" ) == 0 ) {
        decryptmode( argc, argv ); // decryptmode doesn't use any flags, so we can just pass it in.
    } else
    /**
     * Asymmetric mode
     */
    if ( argv[1][0] == 'a' || strcmp( argv[1], "--asymmetric" ) == 0 ) {
        asymmetricmode( argc, argv ); // asymmetricmode doesn't use any flags, so we can just pass it in.
    } else {
        showhelp( argv[0], argv[1] );
        return 1;
    }
    return 0;
}
